import os
import time
# SUMO 带交通仿真和视频创建的 转换 API
from commonroad.visualization.mp_renderer import MPRenderer
import numpy as np
from crdesigner.ui.gui.mwindow.animated_viewer_wrapper.gui_sumo_simulation import SUMO_AVAILABLE
from commonroad.common.file_writer import CommonRoadFileWriter, OverwriteExistingFile
from commonroad.common.file_reader import CommonRoadFileReader

from crdesigner.map_conversion.sumo_map.config import SumoConfig
from crdesigner.map_conversion.sumo_map.cr2sumo.converter import CR2SumoMapConverter
from sumocr.interface.sumo_simulation import SumoSimulation
from sumocr.visualization.video import create_video
# ROOT_DIR = os.path.join(os.getcwd(), '..')

# 设置 commonroad 文件
input_file = 'scenario/commonroad/QiCheChengDiTieZhan_sumo_to_cr.xml'
simu_out_file = 'scenario/commonroad/QiCheChengDiTieZhan_sumo_to_cr_simu.xml'

# 构建 scenario, planning_problem
scenario, planning_problem = CommonRoadFileReader(input_file).open()

# 转换场景到中心 找到当前场景的中心点
centroid = np.mean(np.concatenate(
    [la.center_vertices for la in scenario.lanelet_network.lanelets]),axis=0)
scenario.translate_rotate(-centroid, 0)
planning_problem.translate_rotate(-centroid, 0)

config = SumoConfig.from_scenario_name('QiCheChengDiTieZhan')

# 将 cr 转换为 sumo net
wapper = CR2SumoMapConverter(scenario, config)
# wapper.convert_to_net_file(os.path.join(ROOT_DIR, 'scenario/sumo/QiCheChengDiTieZhan_cr_to_sumo.net.xml')) # TODO 没有这个函数，暂时跳过

# 自动生成交通信号灯
tls_lanelet_id = 43513
traffic_light_system_generated = wapper.auto_generate_traffic_light_system(tls_lanelet_id)
print(f"在 {tls_lanelet_id} 生成了交通信号灯 {traffic_light_system_generated}")


# 在场景生成完交通灯后，绘制场景
rnd = MPRenderer()
wapper.lanelet_network.draw(rnd)
rnd.render(show=True)

# 写入修改后的信号灯
pass

# 执行仿真
simulation = SumoSimulation()
simulation.initialize(config, wapper)

for t in range(config.simulation_steps):
    simulation.simulate_step()
    # time.sleep(0.1)
    print(t)

simulation.stop()

# 保存文件
simulated_scenario = simulation.commonroad_scenarios_all_time_steps()
CommonRoadFileWriter(simulated_scenario,
                     planning_problem,
                     author=scenario.author,
                     affiliation=scenario.affiliation,
                     source=scenario.source,
                     tags=scenario.tags,
                     location=scenario.location).write_scenario_to_file(
    os.path.join(simu_out_file),
    overwrite_existing_file=OverwriteExistingFile.ALWAYS)

print("creating video (this may take some time)")
create_video(simulation.commonroad_scenarios_all_time_steps(),
             'scenario/commonroad', trajectory_pred=simulation.ego_vehicles)

